#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
推荐系统测试

测试推荐系统功能
"""

import sys
import os
from loguru import logger

# 添加项目根目录到路径
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# 在导入任何库之前强制设置环境变量
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
os.environ["HUGGINGFACE_CO_URL_HOME"] = "https://hf-mirror.com"
os.environ["HUGGINGFACE_HUB_CACHE"] = os.path.expanduser("~/.cache/huggingface")
os.environ["SENTENCE_TRANSFORMERS_HOME"] = os.path.expanduser("~/.cache/sentence_transformers")

# 强制导入配置，确保镜像环境变量已设置
import questionretrieval.config

logger.info("🌐 使用 HuggingFace 国内镜像配置")
logger.info(f"📁 缓存目录: {questionretrieval.config.HF_HUB_CACHE}")
logger.info(f"🔗 镜像地址: {questionretrieval.config.HF_MIRROR_ENDPOINT}")

import pytest
from questionretrieval import load_questions

# 测试数据
TEST_QUESTIONS = [
    "大伙房水库的主管部门是哪个单位？",
    "大伙房水库总库容/防洪库容/兴利库容/死库容？",
    "大伙房水库调度目标和任务是什么？",
    "大伙房水库防洪调度的核心任务是什么？",
    "大伙房水库供水调度的核心任务是什么？"
]

TEST_QUERY = "大伙房水库的管理和调度"

def get_test_questions():
    """
    获取测试问题
    """
    try:
        questions = load_questions()
        if len(questions) > 0:
            return questions
    except:
        pass
    return TEST_QUESTIONS

def test_sentence_transformer_recommend():
    """
    测试 SentenceTransformer 检索器推荐
    """
    logger.info("\n=== 测试 SentenceTransformer 推荐 ===")
    
    try:
        from questionretrieval import SentenceTransformerRetriever
        retriever = SentenceTransformerRetriever()
        questions = get_test_questions()
        
        retriever.load_candidates(questions)
        result = retriever.recommend(TEST_QUERY)
        
        # 处理 find_similar_questions 返回的字典列表格式
        if isinstance(result, list) and result:
            logger.info(f"✓ 推荐成功，返回 {len(result)} 个结果")
            if isinstance(result[0], dict) and 'question' in result[0] and 'similarity' in result[0]:
                for i, rec in enumerate(result, 1):
                    logger.info(f"  {i}. {rec['question']} (相似度: {rec['similarity']:.3f})")
            else:
                for i, rec in enumerate(result, 1):
                    logger.info(f"  {i}. {rec}")
        elif isinstance(result, tuple):
            if len(result) == 3:
                recommended_questions, scores, similarities = result
                logger.info(f"✓ 推荐成功，返回 {len(recommended_questions)} 个结果")
                for i, (q, s) in enumerate(zip(recommended_questions, similarities), 1):
                    logger.info(f"  {i}. {q} (相似度: {s:.3f})")
            elif len(result) == 2:
                recommended_questions, scores = result
                logger.info(f"✓ 推荐成功，返回 {len(recommended_questions)} 个结果")
                for i, (q, s) in enumerate(zip(recommended_questions, scores), 1):
                    logger.info(f"  {i}. {q} (分数: {s:.3f})")
        else:
            logger.info(f"✓ 推荐成功，返回结果: {result}")
        
        return True
    except ImportError:
        logger.warning("✗ SentenceTransformer 不可用")
        return False
    except Exception as e:
        logger.error(f"✗ 推荐失败: {e}")
        return False

def test_bert_chinese_recommend():
    """
    测试 BERT 中文检索器推荐
    """
    logger.info("\n=== 测试 BERT 中文检索器推荐 ===")
    
    try:
        from questionretrieval import BertChineseRetriever
        retriever = BertChineseRetriever()
        questions = get_test_questions()
        
        retriever.load_candidates(questions)
        result = retriever.find_similar_questions(TEST_QUERY)
        
        # 处理 find_similar_questions 返回的字典列表格式
        if isinstance(result, list) and result:
            logger.info(f"✓ 推荐成功，返回 {len(result)} 个结果")
            if isinstance(result[0], dict) and 'question' in result[0] and 'similarity' in result[0]:
                for i, rec in enumerate(result, 1):
                    logger.info(f"  {i}. {rec['question']} (相似度: {rec['similarity']:.3f})")
            else:
                for i, rec in enumerate(result, 1):
                    logger.info(f"  {i}. {rec}")
        elif isinstance(result, tuple):
            if len(result) == 3:
                recommended_questions, scores, similarities = result
                logger.info(f"✓ 推荐成功，返回 {len(recommended_questions)} 个结果")
                for i, (q, s) in enumerate(zip(recommended_questions, similarities), 1):
                    logger.info(f"  {i}. {q} (相似度: {s:.3f})")
            elif len(result) == 2:
                recommended_questions, scores = result
                logger.info(f"✓ 推荐成功，返回 {len(recommended_questions)} 个结果")
                for i, (q, s) in enumerate(zip(recommended_questions, scores), 1):
                    logger.info(f"  {i}. {q} (分数: {s:.3f})")
        else:
            logger.info(f"✓ 推荐成功，返回结果: {result}")
        
        return True
    except ImportError:
        logger.warning("✗ BertChineseRetriever 不可用")
        return False
    except Exception as e:
        logger.error(f"✗ 推荐失败: {e}")
        return False

def test_bert_similarity_recommend():
    """
    测试 BERT 相似度检索器推荐
    """
    logger.info("\n=== 测试 BERT 相似度检索器推荐 ===")
    
    try:
        from questionretrieval import BertSimilarityRetriever
        retriever = BertSimilarityRetriever()
        questions = get_test_questions()
        
        retriever.load_candidates(questions)
        result = retriever.find_similar_questions(TEST_QUERY)
        
        # 处理 find_similar_questions 返回的字典列表格式
        if isinstance(result, list) and result:
            logger.info(f"✓ 推荐成功，返回 {len(result)} 个结果")
            if isinstance(result[0], dict) and 'question' in result[0] and 'similarity' in result[0]:
                for i, rec in enumerate(result, 1):
                    logger.info(f"  {i}. {rec['question']} (相似度: {rec['similarity']:.3f})")
            else:
                for i, rec in enumerate(result, 1):
                    logger.info(f"  {i}. {rec}")
        elif isinstance(result, tuple):
            if len(result) == 3:
                recommended_questions, scores, similarities = result
                logger.info(f"✓ 推荐成功，返回 {len(recommended_questions)} 个结果")
                for i, (q, s) in enumerate(zip(recommended_questions, similarities), 1):
                    logger.info(f"  {i}. {q} (相似度: {s:.3f})")
            elif len(result) == 2:
                recommended_questions, scores = result
                logger.info(f"✓ 推荐成功，返回 {len(recommended_questions)} 个结果")
                for i, (q, s) in enumerate(zip(recommended_questions, scores), 1):
                    logger.info(f"  {i}. {q} (分数: {s:.3f})")
        else:
            logger.info(f"✓ 推荐成功，返回结果: {result}")
        
        return True
    except ImportError:
        logger.warning("✗ BertSimilarityRetriever 不可用")
        return False
    except Exception as e:
        logger.error(f"✗ 推荐失败: {e}")
        return False

def test_dbscan_clustering_recommend():
    """
    测试 DBSCAN 聚类检索器推荐
    """
    logger.info("\n=== 测试 DBSCAN 聚类检索器推荐 ===")
    
    try:
        from questionretrieval import DBSCANClusteringRetriever
        retriever = DBSCANClusteringRetriever()
        questions = get_test_questions()
        
        retriever.load_candidates(questions)
        result = retriever.recommend(TEST_QUERY)
        
        # 处理 recommend 返回的元组格式
        if isinstance(result, list) and result:
            logger.info(f"✓ 推荐成功，返回 {len(result)} 个结果")
            for i, rec in enumerate(result, 1):
                if isinstance(rec, tuple) and len(rec) >= 2:
                    question, score = rec[0], rec[1]
                    logger.info(f"  {i}. {question} (分数: {score:.3f})")
                else:
                    logger.info(f"  {i}. {rec}")
        elif isinstance(result, tuple):
            if len(result) >= 2:
                recommended_questions, scores = result[0], result[1]
                logger.info(f"✓ 推荐成功，返回 {len(recommended_questions)} 个结果")
                for i, (q, s) in enumerate(zip(recommended_questions, scores), 1):
                    logger.info(f"  {i}. {q} (分数: {s:.3f})")
        else:
            logger.info(f"✓ 推荐成功，返回结果: {result}")
        
        return True
    except ImportError:
        logger.warning("✗ DBSCANClusteringRetriever 不可用")
        return False
    except Exception as e:
        logger.error(f"✗ 推荐失败: {e}")
        return False

def test_hierarchical_clustering_recommend():
    """
    测试层级聚类检索器推荐
    """
    logger.info("\n=== 测试层级聚类检索器推荐 ===")
    
    try:
        from questionretrieval import HierarchicalClusteringRetriever
        retriever = HierarchicalClusteringRetriever()
        questions = get_test_questions()
        
        retriever.load_candidates(questions)
        result = retriever.recommend(TEST_QUERY)
        
        # 处理 recommend 返回的元组格式
        if isinstance(result, list) and result:
            logger.info(f"✓ 推荐成功，返回 {len(result)} 个结果")
            for i, rec in enumerate(result, 1):
                if isinstance(rec, tuple) and len(rec) >= 2:
                    question, score = rec[0], rec[1]
                    logger.info(f"  {i}. {question} (分数: {score:.3f})")
                else:
                    logger.info(f"  {i}. {rec}")
        elif isinstance(result, tuple):
            if len(result) >= 2:
                recommended_questions, scores = result[0], result[1]
                logger.info(f"✓ 推荐成功，返回 {len(recommended_questions)} 个结果")
                for i, (q, s) in enumerate(zip(recommended_questions, scores), 1):
                    logger.info(f"  {i}. {q} (分数: {s:.3f})")
        else:
            logger.info(f"✓ 推荐成功，返回结果: {result}")
        
        return True
    except ImportError:
        logger.warning("✗ HierarchicalClusteringRetriever 不可用")
        return False
    except Exception as e:
        logger.error(f"✗ 推荐失败: {e}")
        return False

def test_all_retrievers_recommend():
    """
    测试所有检索器推荐功能
    """
    logger.info("\n" + "="*50)
    logger.info("检索器推荐功能测试")
    logger.info("="*50)
    
    questions = get_test_questions()
    logger.info(f"使用 {len(questions)} 个候选问题")
    logger.info(f"查询问题: {TEST_QUERY}")
    
    results = {
        'SentenceTransformer': test_sentence_transformer_recommend(),
        'BertChinese': test_bert_chinese_recommend(),
        'BertSimilarity': test_bert_similarity_recommend(),
        'DBSCANClustering': test_dbscan_clustering_recommend(),
        'HierarchicalClustering': test_hierarchical_clustering_recommend()
    }
    
    logger.info("\n" + "-"*40)
    logger.info("推荐功能测试结果")
    logger.info("-"*40)
    
    success_count = 0
    for name, success in results.items():
        status = "✓ 成功" if success else "✗ 失败"
        logger.info(f"{name:20} : {status}")
        if success:
            success_count += 1
    
    logger.info(f"\n总计: {success_count}/{len(results)} 个检索器可以正常推荐")
    
    if success_count > 0:
        logger.info("\n🎉 恭喜! 有检索器可以正常工作")
        logger.info("💡 提示: 可以在实际项目中使用这些检索器")
    else:
        logger.warning("\n⚠️  没有检索器可以正常推荐")
        logger.info("解决方案:")
        logger.info("1. 检查依赖安装: pip install transformers sentence-transformers torch")
        logger.info("2. 检查网络连接 (首次使用需要下载模型)")
        logger.info("3. 检查问题加载是否成功")
    
    return results

if __name__ == "__main__":
    test_all_retrievers_recommend()